from src.core import *
from src.utils import *
import os

def choose_label_num1(label):
    if label == 'g0':
        label_num = 1
    elif label == 'g1':
        label_num = 2
    elif label == 'g2':
        label_num = 3
    elif label == 'g3':
        label_num = 4
    else:
        raise ValueError('标签种类错误')
    return label_num

'''
labels ：这里为等级g0-g3的排列组合
'''
def extract_feratures(imgs_folder, method, labels, parameters, ratio, save_folder):
    method_folder = method
    save_txt_path = os.path.join(save_folder, method_folder)
    data = Data()
    ft = Feature()
    for label in labels:
        imgs_path = []
        label_imgs_path = os.path.join(imgs_folder, label)
        data.find_inside_file(imgs_path, label_imgs_path) # 取得图片路径，一次取一种种类
        trained_features, tested_features = ft.extract_in_scale(imgs_path, method, parameters=parameters, ratio=ratio,
                                                                if_shuffle=True) # 提取特征
        # 保存特征文件
        ft.save_txt(trained_features, choose_label_num1(label), os.path.join(save_txt_path, "train", label + ".txt"))
        ft.save_txt(tested_features, choose_label_num1(label), os.path.join(save_txt_path, "test", label + ".txt"))


method = 'HOG'
labels = ['g0','g1','g2','g3']
ratio = 0.5
imgs_folder = r'H:\wangjianlian\data\pretreatment\thrid_generation_augmentation'
save_folder = r'H:\wangjianlian\project\Python\image-pretreatment\resources\features'

# 自设的一些参数
if method == "HOG":
    parameters = [9, (64,64), (2,2), None, False, None, False, True, None]
elif method == "LBP":
    parameters = [16, 2, 'uniform'] # 用uniform模式可以有效地降低特征数量（从几万到几十）
elif method == "GLCM":
    parameters = [[5], [0], 256, True, True] # GLCM

extract_feratures(imgs_folder, method, labels, parameters, ratio, save_folder)